# -*- coding: utf-8 -*-
"""
Created on Mon Dec 17 22:23:03 2012

@author: maxim
"""
import aircraft
from numpy import array, dot, arange, zeros, vstack, argmin, flipud, transpose
from math import sin, cos, radians, tan, pi
from matplotlib.pyplot import figure,plot,hold,xlabel,ylabel,show,grid,axis
from mpl_toolkits.mplot3d import Axes3D
from geometry import ellipse2D

class DisplayAircraft:
    """
    Plots aicraft configuration using matplotlib library.
    All the geometry can be separated into several types:
        
        1. fuselage - fuselage or other bodies with similar geometry (nacelle)
        2. wing - wing with specified airfoils and other specific wing geometry
        3. mass - all the mass components are displayed as points with labels
        4. landing gear - by given position, diameter and width landing gear 
           are displayed as 3D torus objects
    """
    
    def __init__(self):
        self.Xmsh     = list()
        self.Ymsh     = list()
        self.Zmsh     = list()
        self.Xpts     = list()
        self.Ypts     = list()
        self.Zpts     = list()
        self.color    = list()
        self.opacity  = list()
        self.text     = list()
        self.pointColor = list()
        self.axisLim  = array([-5.,5.])
    
    def add_wing(self,wing,color='b',opacity=0.8,numPts=20):
        """
        adds new wing to aircraft plot. Default wing will be created with 
        airfoil section in XZ plane, span direction +Y and symmetry wrt XZ plane.
        If vStab=True airfoil section is in XY plane, span direction +Z, no symmetry.
        
        :param wing: aircraft.wing object
        :param color: color of the wing (ref: matplotlib)
        :param opacity: opacity of the wing
        :param integer numPts: number of points on airfoil surface.
        """
        nSeg = len(wing.segmentSpans)
        refPt = wing.aapex
        refPt[1] += wing.segmentSpans[0]
        for i in range(nSeg):
            airfoil = wing.airfoil[i].redim(numPts,overwrite=False)
            section = airfoil*wing.chords[i]
            axis1 = [wing.chords[i]*0.25,0]
            if i==0:
                angle = wing.incidence
                section = self.rotate2D(section,axis1,angle)
                wingX = section[:,0] + refPt[0]
                wingY = zeros(numPts)
                wingZ = section[:,1] + refPt[2]
            else:
                angle = wing.incidence + wing.twistAngles[i-1]
                section = self.rotate2D(section,axis1,angle)
                refPt[0] += wing.segmentOffsets[i-1]
                refPt[1] += wing.segmentSpans[i]
                refPt[2] += wing.segmentSpans[i]*tan(radians(wing.segmentDihedrals[i-1]))
                secX = section[:,0] + refPt[0]
                secY = zeros(numPts)  + refPt[1]
                secZ = section[:,1] + refPt[2]
                wingX = vstack([secX,wingX,secX])
                wingY = vstack([-secY,wingY,secY])
                wingZ = vstack([secZ,wingZ,secZ])
        self.Xmsh.append(wingX)
        self.Ymsh.append(wingY)
        self.Zmsh.append(wingZ)
        self.color.append(color)
        self.opacity.append(opacity)
    
    def add_vstab(self,vstab,color='y',opacity=0.8,numPts=20):
        nSeg = len(vstab.segmentSpans)
        refPt = vstab.aapex
        refPt[2] += vstab.segmentSpans[0]
        for i in range(nSeg):
            airfoil = vstab.airfoil[i].redim(numPts,overwrite=False)
            section = airfoil*vstab.chords[i]
            axis1 = [vstab.chords[i]*0.25,0.0]
            if i==0:
                angle = vstab.incidence
                section = self.rotate2D(section,axis1,angle)
                vstabX = section[:,0] + refPt[0]
                vstabY = section[:,1] + refPt[1]
                vstabZ = zeros(numPts)
            else:
                angle = vstab.incidence + vstab.twistAngles[i-1]
                section = self.rotate2D(section,axis1,angle)
                refPt[0] += vstab.segmentOffsets[i-1]
                refPt[1] += vstab.segmentSpans[i]*tan(radians(vstab.segmentDihedrals[i-1]))
                refPt[2] += vstab.segmentSpans[i]
                secX = section[:,0] + refPt[0]
                secY = section[:,1] + refPt[1]
                secZ = zeros(numPts)  + refPt[2]
                vstabX = vstack([vstabX,secX])
                vstabY = vstack([vstabY,secY])
                vstabZ = vstack([vstabZ,secZ])
        self.Xmsh.append(vstabX)
        self.Ymsh.append(vstabY)
        self.Zmsh.append(vstabZ)
        self.color.append(color)
        self.opacity.append(opacity)
    
    def add_fuselage(self,fus,color='y',opacity=0.8,nCsPts=10):
        """
        Append fuselage to plot.
        
        :ivar fus: fuselage object from aircraft.py module
        :ivar color: color of fuselage
        :ivar opacity: set opacity of the fuselage
        :ivar nCsPts: number of cross section points
        
        Ellipse is used for fuselage cross section
        """
        nPts = len(fus.topProfile_X)
        sideCurve = array([fus.topProfile_X, fus.topProfile_Y,zeros(nPts)])
        sideCurve = transpose(sideCurve) #TODO: modify to read fuselage.profileSide...
        nPts = len(fus.sideProfile_X)
        sideViewCurve = array([fus.sideProfile_X, zeros(nPts), fus.sideProfile_Y])
        sideViewCurve = transpose(sideViewCurve)
        upCurve, loCurve = self._separate_fus_points(sideViewCurve)
        
        for i,upPt in enumerate(upCurve):
            cs = self.ellipse_cs(upPt,sideCurve[i],loCurve[i],nCsPts)
            if i==0:
                fusX,fusY,fusZ = cs[:,0], cs[:,1], cs[:,2]
            else:
                fusX = vstack([fusX,cs[:,0]])
                fusY = vstack([fusY,cs[:,1]])
                fusZ = vstack([fusZ,cs[:,2]])
        self.color.append(color)
        self.opacity.append(opacity)
        self.Xmsh.append(fusX)
        self.Ymsh.append(fusY)
        self.Zmsh.append(fusZ)

    def add_landing_gear(self,lg,color='y',opacity=1.0):
        """
        :ivar lg: aircraft.landingGear object
        
        adds landing gear to aircraft plot. Landing gear is displayed as 3D torus
        with dimmension based on diameter and width of the tire
        """
        nLG = len(lg.tireDiameter)
        for i in range(nLG):
            r = lg.tireWidth[i]/2
            R = lg.tireDiameter[i]/2-r
            center = [lg.groundContact_X[i],lg.groundContact_Y[i],lg.groundContact_Z[i]+R+r]
            lgX,lgY,lgZ = self._torus3D(R,r,center)
            self.Xmsh.append(lgX)
            self.Ymsh.append(lgY)
            self.Zmsh.append(lgZ)
            self.color.append(color)
            self.opacity.append(opacity)
    
    def add_propeller(self,prop,color='k',opacity=1.0):
        pass
    
    def add_mass(self,mass,color='r',marker='o'):
        """
        adds point mass to the plot with label: mass of the component
        
        :ivar mass: aircraft.massList object
        :ivar color: color of point mass marker
        :ivar marker: shape of the marker
        """
        massX, massY, massZ,text = list(),list(),list(),list()
        for item in mass.total.items:
            massX.append(item.coords[0])
            massY.append(item.coords[1])
            massZ.append(item.coords[2])
            text.append('%.1fkg'%item.mass)
        self.Xpts.append(massX)
        self.Ypts.append(massY)
        self.Zpts.append(massZ)
        self.pointColor.append(color+marker)
        self.text.append(text)

    def add_point_data(self,points,text,color='g',marker='o'):
        self.Xpts.append(points[:,0])
        self.Ypts.append(points[:,1])
        self.Zpts.append(points[:,2])
        self.text.append(text)
        self.pointColor.append(color+marker)
    
    def plot(self,mod='surf'):
        fig = figure(1)
        hold(True)
        ax = Axes3D(fig)
        #ax = fig.gca(projection='3d')
        for i,mshX in enumerate(self.Xmsh):
            mshY = self.Ymsh[i]
            mshZ = self.Zmsh[i]
            color = self.color[i]
            transp = self.opacity[i]
            if mod=='wire':
                ax.plot_wireframe(mshX,mshY,mshZ,color=color,alpha=transp,
                               linewidth=0.5,rstride=1, cstride=1)
            else:
                ax.plot_surface(mshX,mshY,mshZ,color=color,alpha=transp,
                               linewidth=0.0,rstride=1, cstride=1)
        for i,xpts in enumerate(self.Xpts):
            ypts = self.Ypts[i]
            zpts = self.Zpts[i]
            txt  = self.text[i]
            col  = self.pointColor[i]
            ax.plot(xpts,ypts,zpts,col)
            for j,x in enumerate(xpts):
                ax.text(x,ypts[j],zpts[j],txt[j])

        ax.set_aspect('equal')
        #ax.set_xlim(-1,6)
        #ax.set_xlim(self.axisLim[0],self.axisLim[1])
        ax.set_xlim(-2,8)
        ax.set_ylim(self.axisLim[0],self.axisLim[1])
        ax.set_zlim(self.axisLim[0],self.axisLim[1])
        show()
    
    def _torus3D(self,R,r,center,nPts=[10,10]):
        pi2 = 2*pi
        step1 = pi2/(nPts[0]-1)
        step2 = pi2/(nPts[1]-1)
        _phi = arange(0,pi2+step1,step1)
        _ksi = arange(0,pi2+step2,step2)
        X = zeros([nPts[1],nPts[0]])
        Y = zeros([nPts[1],nPts[0]])
        Z = zeros([nPts[1],nPts[0]])
        for i,ksi in enumerate(_ksi):
            for j,phi in enumerate(_phi):
                X[i,j] = (R + r*cos(phi))*cos(ksi)
                Z[i,j] = (R + r*cos(phi))*sin(ksi)
                Y[i,j] = r*sin(phi)
        return X+center[0],Y+center[1],Z+center[2]

    def rotate2D(self,curve, axis, angle):
        """
        rotates 2D curve around 2D axis point to a given angle
        
        :ivar array curve: array of curve points
        :ivar array axis: rotation axis in format array([x,y])
        :ivar float angle: rotation anlge in degrees
        
        :math:`x_{new} = cos(\theta)x - sin(\theta)y`
        
        :math:`y_{new} = sin(\theta)x + cos(\theta)y`
        """
        angle = radians(angle)
        rotMatrix = array([[cos(angle), -sin(angle)],[sin(angle), cos(angle)]])
        return dot( (curve-axis), rotMatrix) + axis
    
    def ellipse_cs(self,ptUp,ptSide,ptLo,numPts=7):
        """
        Creates ellipse cross section for body-like objects.
        
        :ivar 3Darray ptUp: upper point of cross section
        :ivar 3Darray ptSide: side point of cross section
        :ivar 3Darray ptLo: lower point of cross section
        
        note: all points should have same X position.
        :returns: array of cross section points
        """
        csUp = zeros([numPts,3])
        csLo = zeros([numPts,3])
        center = array([ptLo[1],ptSide[2]])
        csLo[:,1:3] = ellipse2D(ptLo[1:3], ptSide[1:3],center,-180,0, numPts)
        csUp[:,1:3] = ellipse2D(ptSide[1:3], ptUp[1:3],center, 0, 180, numPts)
        csLo[:,0] = ptUp[0]
        csUp[:,0] = ptUp[0]
        cs = vstack([csLo, csUp,csLo[0]])
        return cs

    def _separate_fus_points(self,pts):
        sepPtIdx = argmin(pts[:,0])
        return flipud(pts[:sepPtIdx+1]), pts[sepPtIdx:]
    def add_wing2(self,wing):
        self.Xmsh.append(wing.meshX)
        self.Ymsh.append(wing.meshY)
        self.Zmsh.append(wing.meshZ)
        self.color.append(wing.color)
        self.opacity.append(wing.opacity)

def plot_aircraft(ac):
    acDisp = DisplayAircraft()
    acDisp.add_fuselage(ac.fuselage,color='y',opacity=0.5)
    acDisp.add_wing(ac.wing,color='b',opacity=0.5)
    acDisp.add_wing(ac.hStab,color='b',opacity=0.5)
    acDisp.add_vstab(ac.vStab,color='b',opacity=0.5)
    acDisp.add_mass(ac.mass,'r','o')
    acDisp.add_landing_gear(ac.landingGear,color='k')
    acDisp.plot('surf')

def plot_aircraft2(ac):
    acDisp = DisplayAircraft()
    acDisp.add_fuselage(ac.fuselage,color='y',opacity=0.5)
    acDisp.add_wing2(ac.wing)
    acDisp.add_wing2(ac.hStab)
    acDisp.add_wing2(ac.vStab)
    acDisp.add_mass(ac.mass,'r','o')
    acDisp.add_landing_gear(ac.landingGear,color='k')
    acDisp.plot('surf')

def test_function():
    ac = aircraft.load("sampleInput2")
    plot_aircraft(ac) 

if __name__=="__main__":
    test_function()